
#include <vector>
#include <iostream>
#include <thread>

typedef struct{
    int start;
    int end;
    int dim;
    int m;
    int n;
    float *a;
    float *b;
}Nos;

void sum_t(Nos *p){
    
    for (int idx = p->start; idx < p->end; idx++) {
        if (p->dim == 0){
        for(int k = 0 ; k < p->m; k++){
            p->b[idx] += p->a[k * p->n + idx];
            }
        }else if(p->dim == 1){
            for(int k = 0 ; k < p->n; k++){
                p->b[idx] += p->a[idx * p->n + k];
            }
        }
    }
}

void sum(float *a, float *b, int m, int n, int dim){
    int size;
    if (dim == 0){
        size = n;
    }
    else if(dim == 1){
        size = m;
    }
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Nos> mes(thread_num);
    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Nos{start, end, dim, m, n, a, b};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(sum_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
}

int main(){
    int m = 2;
    int n = 3;
    int dim = 1;
    float a[6] = {1,2,3,4,5,6};
    float b[6] = {0};
    sum(a,b,m,n,dim);
    for (int i = 0; i < 6; i++) {
        std::cout << b[i] << " ";
    }
    std::cout << std::endl;
}